1004E - Sonya and Ice Cream - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
 
#define FOR(i,a,b) for (int i = (a); i < (b); ++i)
#define F0R(i,a) FOR(i,0,a)
#define ROF(i,a,b) for (int i = (b)-1; i >= (a); --i)
#define R0F(i,a) ROF(i,0,a)
 
using ll = long long;
using ld = long double;
ll INF = LLONG_MAX;
 
using vi = vector<int>;
using vll = vector<ll>;
using pii = pair<int, int>;
 
namespace output {
    void pr(int x) { cout << x; }
    void pr(long x) { cout << x; }
    void pr(ll x) { cout << x; }
    void pr(unsigned x) { cout << x; }
    void pr(unsigned long x) { cout << x; }
    void pr(unsigned long long x) { cout << x; }
    void pr(float x) { cout << x; }
    void pr(double x) { cout << x; }
    void pr(ld x) { cout << x; }
    void pr(char x) { cout << x; }
    void pr(const char* x) { cout << x; }
    void pr(const string& x) { cout << x; }
    void pr(bool x) { pr(x ? "true" : "false"); }
    template<class T> void pr(const complex<T>& x) { cout << x; }
    template<class T1, class T2> void pr(const pair<T1,T2>& x);
    template<class T> void pr(const T& x);
    template<class T, class... Ts> void pr(const T& t, const Ts&... ts) { 
        pr(t); pr(ts...); 
    }
    template<class T1, class T2> void pr(const pair<T1,T2>& x) { 
        pr("{",x.f,", ",x.s,"}");
    }
    template<class T> void pr(const T& x) { 
        pr("{"); // const iterator needed for vector<bool>
        bool fst = 1; for (const auto& a: x) pr(!fst?", ":"",a), fst = 0; 
        pr("}");
    }
    void print() { pr("\n"); } // print w/ spaces
    template<class T, class... Ts> void print(const T& t, const Ts&... ts) { 
        pr(t); if (sizeof...(ts)) pr(" "); print(ts...); 
    }
}
 
using namespace output;

int N, K; 
vector<vector<pair<int, ll>>> graph;
vector<ll> dist;
vector<ll> pre;
vector<ll> height;

void dfs(int cur) {
    for (auto i : graph[cur]) if (i.first != pre[cur]) {
        pre[i.first] = cur;
        dist[i.first] = dist[cur] + i.second;
        dfs(i.first);
        height[cur] = max(height[cur], height[i.first] + i.second);
    }
}

void genDist(int cur) {
    dist = vector<ll>(N);
    pre = vector<ll> (N);
    height = vector<ll> (N);
    pre[cur] = -1;
    dfs(cur);
}

pair<int, int> treeDiameter() {
    genDist(0);
    int best = 0; F0R(i, N) if (dist[i] > dist[best]) best = i;
    genDist(best);
    int start = best;
    F0R(i, N) if (dist[i] > dist[best]) best = i;
    return {start, best};
}

int main() {
    ios::sync_with_stdio(0), cin.tie(0), cout.tie(0);
    cin >> N >> K;
    graph.resize(N);
    F0R(i, N-1) {
        int a, b, c; cin >> a >> b >> c; --a; --b;
        graph[a].push_back({b, c});
        graph[b].push_back({a, c});
    }

    pair<int, int> diam = treeDiameter();
    int s = diam.first;
    int e = diam.second;
    ll diamLen = dist[e];

    vi diamPath = {};
    //print("DFS", s, "->", e);
    int cur = e;
    while (e != -1) {
        diamPath.push_back(e);
        e = pre[e];
    }
    reverse(diamPath.begin(), diamPath.end());

    ll ans = INF;
    if (K < diamPath.size()) {
        for (int i = 0; i + K - 1 < diamPath.size(); ++i) {
            ans = min(ans, max(dist[diamPath[i]], diamLen - dist[diamPath[i+K-1]]));
        }
    } else {
        ans = 0;
    }
    //print("HEIGHTS", height);
    for (int i = 0; i < diamPath.size(); ++i) {
        for (pii e : graph[diamPath[i]]) if ((i == 0 || diamPath[i-1] != e.first) && (i == diamPath.size() - 1 || diamPath[i+1] != e.first)) {
            ans = max(ans, height[e.first] + e.second);
        }
    }
    print(ans);
}